Quench of the quantum Ising model starting from the ground state.

# pylint: disable=too-many-arguments
# pylint: disable=invalid-name

import os

# pylint: disable-next=no-member, no-name-in-module
import os.path

import matplotlib.pyplot as plt
import numpy as np

import qtealeaves as qtl
from qtealeaves.models import get_quantum_ising_1d

# Keys are L, g_init, tau,  symmetry sector, value is final energy at g_final=0
ref_osmps = {
    (8, 2.0, 10, None) : -6.88725731743711,
    (8, 2.0, 10, 0) : None,
    (8, 2.0, 10, 1) : None
}

# pylint: disable-next=too-many-locals
def main(
        tn_type=5,
        input_folder=None,
        output_folder=None,
        timesteps=200,
        plot=True
        ):
    """
    Main method for quenching the quantum Ising model starting in a
    ground state.

    **Arguments**

    tn_type : int, optional
        Choose 5 for python-TTN, 6 for python-MPS.
        Default to 5.

    input_folder : str | None, optional
        Input folder. Default to None.

    output_folder : str | None, optional
        Output folder. Default to None.

    timesteps : int, optional
        Number of timesteps. Default to 160.

    plot : bool, optional
        If True, plot the results at the end of the example. Default tot True
    """
    if input_folder is None:
        input_folder = lambda params : 'QI1dDyn/input_L%d'%(params['L'])
    if output_folder is None:
        output_folder = lambda params : 'QI1dDyn/output_L%d'%(params['L'])

    model, my_ops = get_quantum_ising_1d()

    my_conv = qtl.convergence_parameters.TNConvergenceParameters(max_iter=7, max_bond_dimension=20)
    my_obs = qtl.observables.TNObservables()

    # Dynamics
    nn_meas = 10
    quench = qtl.DynamicsQuench('t_grid', measurement_period=nn_meas,
                                time_evolution_mode=1)
    quench['g'] = lambda tt, params : 2.0 - 2.0 * (tt / 10.0)

    simulation = qtl.QuantumGreenTeaSimulation(model, my_ops, my_conv, my_obs,
                                    tn_type=tn_type,
                                    folder_name_input=input_folder,
                                               folder_name_output=output_folder,
                                               store_checkpoints=False,
    )

    params = []

    params.append({
        'L' : 8,
        'J' : 1.0,
        'g' : 2.0,
        't_grid' : [0.05] * timesteps,
        'Quenches' : [quench],
        "exclude_from_hash" : ["Quenches"],
    })

    simulation.run(params, delete_existing_folder=True)

    for elem in params:
        obs_quench_1 = simulation.get_dynamic_obs(elem)[0]
        energies = [entry['energy'] for entry in obs_quench_1]

        if plot:
            fig = plt.figure()
            ax1 = fig.add_subplot(111)
            ax1.plot(np.cumsum(elem['Quenches'][0].get_dt_grid(elem))[::nn_meas],
                    energies, 'b-')
            ax1.set_xlabel('Time t')
            ax1.set_ylabel('Energy E')
            # pylint: disable-next=no-member
            plt.savefig(os.path.join(output_folder(elem), 'dyn_energy.pdf'))

    print(
        f"\nExample `{__file__}` ran successfully; "
        + "pdf-figures are saved to QI1dDyn; "
        + "no asserts are implemented for this example."
    )

    return


if(__name__ == '__main__'):
    main()